import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import griddata
from scipy.spatial.distance import cdist
from scipy.interpolate import interp1d
import tkinter as tk
from tkinter import filedialog

DISTANCE = 2  # 设置距离阈值，单位与坐标相同


class ContourPlotter:

    def __init__(self):
        self.df = None
        self.x = None
        self.y = None
        self.z = None
        self.Xi = None
        self.Yi = None
        self.Zi = None
        self.fig = None
        self.ax = None
        self.ax_section = None
        self.lines = []
        self.nearest_point_markers = []
        self.section_lines = []
        self.drawn_lines = []
        self.l_key_pressed = False

        # 设置中文字体支持
        plt.rcParams['font.sans-serif'] = [
            'SimHei', 'FangSong', 'Arial Unicode MS'
        ]
        plt.rcParams['axes.unicode_minus'] = False

    def load_data(self):
        # 创建一个隐藏的Tk根窗口用于文件选择
        root = tk.Tk()
        root.withdraw()  # 隐藏主窗口

        # 弹出文件选择对话框让用户选择CSV文件
        file_path = filedialog.askopenfilename(title="请选择要导入的CSV文件",
                                               filetypes=[
                                                   ("CSV files", "*.csv"),
                                                   ("All files", "*.*")
                                               ])

        # 如果用户没有选择文件则退出
        if not file_path:
            print("未选择文件，程序退出。")
            # exit()
            return

        # 加载CSV数据
        self.df = pd.read_csv(file_path)

        # 提取坐标和高程数据
        self.x = self.df.iloc[:, 2].values  # X坐标
        self.y = self.df.iloc[:, 3].values  # Y坐标
        self.z = self.df.iloc[:, 4].values  # Z坐标(高程)

    def interpolate_data(self):
        # 创建更密集的网格点用于插值
        xi = np.linspace(self.x.min(), self.x.max(), 200)
        yi = np.linspace(self.y.min(), self.y.max(), 200)
        self.Xi, self.Yi = np.meshgrid(xi, yi)

        # 使用griddata进行插值，只在数据点凸包内插值，外部区域设为NaN
        self.Zi = griddata((self.x, self.y),
                           self.z, (self.Xi, self.Yi),
                           method='cubic',
                           fill_value=np.nan)

        # 创建一个遮罩，隐藏距离数据点过远的区域
        # 计算每个网格点到最近数据点的距离
        grid_points = np.column_stack([self.Xi.ravel(), self.Yi.ravel()])
        data_points = np.column_stack([self.x, self.y])
        distances = cdist(grid_points, data_points).min(axis=1)
        distance_threshold = 5  # 设置距离阈值，单位与坐标相同
        mask = distances.reshape(self.Xi.shape) > distance_threshold
        self.Zi[mask] = np.nan

    def setup_plot(self):
        # 创建图形和轴
        self.fig, (self.ax, self.ax_section) = plt.subplots(1,
                                                            2,
                                                            figsize=(15, 7))

        # 创建从深蓝色到黄色的渐变色映射
        colors = ['darkblue', 'blue', 'cyan', 'green', 'yellow']
        n_bins = 20  # 与等高线级别数一致
        cmap = plt.cm.colors.LinearSegmentedColormap.from_list('custom',
                                                               colors,
                                                               N=n_bins)

        # 绘制等高线图
        contour_filled = self.ax.contourf(self.Xi,
                                          self.Yi,
                                          self.Zi,
                                          levels=20,
                                          cmap=cmap)
        contour_lines = self.ax.contour(self.Xi,
                                        self.Yi,
                                        self.Zi,
                                        levels=20,
                                        colors='black',
                                        linewidths=0.8,
                                        alpha=0.9)

        # 添加颜色条
        cbar = self.fig.colorbar(contour_filled, ax=self.ax, label='高程 (m)')

        # 添加等高线标签
        self.ax.clabel(contour_lines, inline=True, fontsize=12, fmt='%.1f')

        # 用细线连接原始数据点而不是显示为点
        self.ax.plot(self.x,
                     self.y,
                     color='white',
                     linewidth=0.5,
                     alpha=0.7,
                     label='原始数据连线')

        # 设置图表标题和轴标签
        self.ax.set_title('水下地形图')
        self.ax.set_xlabel('X 坐标')
        self.ax.set_ylabel('Y 坐标')

        # 初始化断面图
        self.ax_section.set_title('断面图')
        self.ax_section.set_xlabel('距离 (m)')
        self.ax_section.set_ylabel('高程 (m)')

        # 添加说明文本
        self.ax.text(
            0.02,
            0.2,
            '交互说明:\n先按住Ctrl键再左键点击两点绘制直线\n右键点击删除最近的线段\n程序将显示直线上最近的数据点和插值点\n并生成断面图',
            transform=self.ax.transAxes,
            fontsize=10,
            verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

    def find_nearest_points(self, x1, y1, x2, y2):
        # 存储最近点信息
        nearest_points = []

        # 查找原始数据点中的最近点
        for i in range(len(self.x)):
            px, py = self.x[i], self.y[i]
            # 计算点到直线的距离
            # 直线方程: (y2-y1)x - (x2-x1)y + (x2*y1-x1*y2) = 0
            A = y2 - y1
            B = x1 - x2
            C = x2 * y1 - x1 * y2

            # 点到直线距离公式
            dist = abs(A * px + B * py + C) / np.sqrt(A * A + B * B)

            # 检查点是否在线段附近（5米范围内）
            if dist < 5:
                nearest_points.append((px, py, self.z[i], dist))

        # 查找插值网格点中的最近点
        # 将网格点展平为一维数组进行处理
        Xi_flat = self.Xi.ravel()
        Yi_flat = self.Yi.ravel()
        Zi_flat = self.Zi.ravel()

        for i in range(len(Xi_flat)):
            px, py, pz = Xi_flat[i], Yi_flat[i], Zi_flat[i]
            # 跳过无效的插值点
            if np.isnan(pz):
                continue

            # 计算点到直线的距离
            A = y2 - y1
            B = x1 - x2
            C = x2 * y1 - x1 * y2

            dist = abs(A * px + B * py + C) / np.sqrt(A * A + B * B)

            # 检查点是否在线段附近（5米范围内）
            if dist < 5:
                nearest_points.append((px, py, pz, dist))

        # 按距离排序
        nearest_points.sort(key=lambda p: p[3])
        return nearest_points

    def plot_cross_section(self, nearest_points, x1, y1, x2, y2):
        # 显示最近点（最多显示10个）
        displayed_points = []
        for i, point in enumerate(nearest_points[:10]):
            px, py, pz, dist = point
            marker = self.ax.plot(
                px,
                py,
                'bo',
                markersize=8,
                label='最近数据点'
                if len(self.nearest_point_markers) == 0 else "")[0]
            self.nearest_point_markers.append(marker)

            # 添加注释
            self.ax.annotate(f'({px:.1f}, {py:.1f}, {pz:.3f})', (px, py),
                             xytext=(10, 10),
                             textcoords='offset points',
                             bbox=dict(boxstyle='round,pad=0.3',
                                       facecolor='yellow',
                                       alpha=0.7),
                             fontsize=9)

            # 收集用于断面图的数据
            # 计算点在线段上的投影位置（距离起点的距离）
            proj_dist = np.sqrt((px - x1)**2 + (py - y1)**2)
            displayed_points.append((proj_dist, pz))

        # 绘制断面图
        if displayed_points:
            # 按距离排序
            displayed_points.sort(key=lambda p: p[0])
            distances, elevations = zip(*displayed_points)

            # 修改：生成5米间隔的平滑曲线
            # 创建5米间隔的采样点
            min_dist = min(distances)
            max_dist = max(distances)
            sample_distances = np.arange(int(min_dist),
                                         int(max_dist) + 1, DISTANCE)  # 5米间隔

            # 使用样条插值创建平滑曲线
            if len(distances) >= 2:  # 至少需要2个点才能插值
                # 使用三次样条插值创建平滑曲线
                f = interp1d(distances,
                             elevations,
                             kind='cubic',
                             fill_value='extrapolate')
                sample_elevations = f(sample_distances)

                # 绘制平滑曲线和采样点
                # 绘制平滑曲线
                smooth_line, = self.ax_section.plot(sample_distances -
                                                    min(sample_distances),
                                                    sample_elevations,
                                                    'b-',
                                                    linewidth=2,
                                                    label='高程线')
                # 绘制1米间隔的采样点
                sample_points, = self.ax_section.plot(sample_distances -
                                                      min(sample_distances),
                                                      sample_elevations,
                                                      'ro',
                                                      markersize=4,
                                                      label=f'{DISTANCE}米间隔')

                # 在每个采样点下方显示高程数值，避免重叠
                prev_text_pos = None  # 记录上一个文本位置
                min_distance = 30  # 文本之间的最小距离（像素）

                for i, (dist, elev) in enumerate(
                        zip(sample_distances, sample_elevations)):
                    # 计算当前文本的位置
                    current_x = dist - min(sample_distances)
                    current_y = elev

                    # 检查是否需要调整垂直位置以避免重叠
                    y_offset = -15  # 默认向下偏移
                    if prev_text_pos is not None:
                        prev_x, prev_y = prev_text_pos
                        # 转换为像素坐标进行距离计算
                        trans = self.ax_section.transData
                        current_pixel = trans.transform((current_x, current_y))
                        prev_pixel = trans.transform((prev_x, prev_y))

                        # 计算像素距离
                        pixel_distance = abs(current_pixel[0] - prev_pixel[0])

                        # 如果水平距离太近，调整垂直位置
                        if pixel_distance < min_distance:
                            y_offset = -30  # 更大的向下偏移

                    self.ax_section.annotate(
                        f'{elev:.2f}', (current_x, current_y),
                        xytext=(0, y_offset),
                        textcoords='offset points',
                        ha='center',
                        va='top',
                        fontsize=8,
                        bbox=dict(boxstyle='round,pad=0.2',
                                  facecolor='white',
                                  alpha=0.7))

                    # 更新上一个文本位置
                    prev_text_pos = (current_x, current_y)

                self.section_lines.extend([smooth_line, sample_points])

                # 添加网格线作为辅助线
                self.ax_section.grid(True, linestyle='--', alpha=0.7)
                self.ax_section.minorticks_on()
                self.ax_section.grid(which='minor', linestyle=':', alpha=0.5)

                # 设置断面图标题
                self.ax_section.set_title(
                    f'断面图 (方位角: {np.degrees(np.arctan2(y2-y1, x2-x1)):.1f}°)')
                self.ax_section.legend()
            else:
                # 如果点太少，直接绘制原始点
                section_line, = self.ax_section.plot(np.array(distances) -
                                                     min(distances),
                                                     elevations,
                                                     'b-o',
                                                     linewidth=2,
                                                     markersize=4)

                # 在每个原始点下方显示高程数值，避免重叠
                prev_text_pos = None
                min_distance = 30

                for i, (dist, elev) in enumerate(zip(distances, elevations)):
                    current_x = dist - min(distances)
                    current_y = elev

                    y_offset = -15
                    if prev_text_pos is not None:
                        prev_x, prev_y = prev_text_pos
                        trans = self.ax_section.transData
                        current_pixel = trans.transform((current_x, current_y))
                        prev_pixel = trans.transform((prev_x, prev_y))

                        pixel_distance = abs(current_pixel[0] - prev_pixel[0])

                        if pixel_distance < min_distance:
                            y_offset = -30

                    self.ax_section.annotate(
                        f'{elev:.2f}', (current_x, current_y),
                        xytext=(0, y_offset),
                        textcoords='offset points',
                        ha='center',
                        va='top',
                        fontsize=8,
                        bbox=dict(boxstyle='round,pad=0.2',
                                  facecolor='white',
                                  alpha=0.7))

                    prev_text_pos = (current_x, current_y)

                self.section_lines.append(section_line)

                # 添加网格线作为辅助线
                self.ax_section.grid(True, linestyle='--', alpha=0.7)
                self.ax_section.minorticks_on()
                self.ax_section.grid(which='minor', linestyle=':', alpha=0.5)

                self.ax_section.set_title(
                    f'断面图 (方位角: {np.degrees(np.arctan2(y2-y1, x2-x1)):.1f}°)')

    def on_click(self, event):
        if event.inaxes != self.ax:
            return

        # 记录点击位置
        # 修改：只有当'Ctrl'键被按下时才处理左键点击
        if event.button == 1 and self.l_key_pressed:  # 左键点击且'Ctrl'键已按下
            if len(self.lines) == 0 or len(self.lines[-1]) == 2:
                # 开始新线段
                self.lines.append([(event.xdata, event.ydata)])
                # 清除之前的最近点标记
                for marker in self.nearest_point_markers:
                    marker.remove()
                self.nearest_point_markers.clear()
                # 清除之前的断面线
                for line in self.section_lines:
                    line.remove()
                self.section_lines.clear()
                # 清除断面图内容
                self.ax_section.clear()
                self.ax_section.set_title('断面图')
                self.ax_section.set_xlabel('距离 (m)')
                self.ax_section.set_ylabel('高程 (m)')
            else:
                # 完成当前线段
                self.lines[-1].append((event.xdata, event.ydata))
                # 绘制线段
                line_x = [self.lines[-1][0][0], self.lines[-1][1][0]]
                line_y = [self.lines[-1][0][1], self.lines[-1][1][1]]
                line_obj, = self.ax.plot(
                    line_x,
                    line_y,
                    'r-',
                    linewidth=2,
                    label='用户绘制线段' if len(self.drawn_lines) == 0 else "")
                self.drawn_lines.append(line_obj)

                # 计算直线方程 Ax + By + C = 0
                x1, y1 = self.lines[-1][0]
                x2, y2 = self.lines[-1][1]

                # 查找最近点
                nearest_points = self.find_nearest_points(x1, y1, x2, y2)

                # 绘制断面图
                self.plot_cross_section(nearest_points, x1, y1, x2, y2)

                # 添加图例
                self.ax.legend()

                # 刷新图形
                self.fig.canvas.draw()

        # 右键点击删除最近的线段
        elif event.button == 3:  # 右键点击
            if self.drawn_lines:
                # 找到最近的线段并删除
                min_dist = float('inf')
                line_to_remove = None
                line_index = -1

                for i, line in enumerate(self.drawn_lines):
                    # 获取线段的坐标
                    line_x, line_y = line.get_data()
                    # 计算点击点到线段中点的距离
                    mid_x = (line_x[0] + line_x[1]) / 2
                    mid_y = (line_y[0] + line_y[1]) / 2
                    dist = np.sqrt((event.xdata - mid_x)**2 +
                                   (event.ydata - mid_y)**2)

                    if dist < min_dist:
                        min_dist = dist
                        line_to_remove = line
                        line_index = i

                # 如果找到了最近的线段并且距离足够近，则删除它
                if line_to_remove and min_dist < 20:  # 20米范围内认为是点击了该线段
                    # 删除线段对象
                    line_to_remove.remove()
                    # 从列表中移除
                    self.drawn_lines.pop(line_index)
                    # 从lines数据中移除对应的坐标
                    self.lines.pop(line_index)

                    # 修改：只删除与这条线段相关的标记和注释
                    # 先清除所有标记和注释
                    for marker in self.nearest_point_markers:
                        marker.remove()
                    self.nearest_point_markers.clear()

                    # 添加：清除高程点注释文本
                    texts_to_remove = []
                    for txt in self.ax.texts:
                        # 检查是否为高程点注释（根据文本内容判断）
                        if txt.get_text().startswith(
                                '(') and txt.get_text().endswith(')'):
                            texts_to_remove.append(txt)

                    for txt in texts_to_remove:
                        txt.remove()

                    # 重新计算并绘制剩余线段的标记点
                    for i, line_coords in enumerate(self.lines):
                        # 重新计算每个线段的最近点
                        x1, y1 = line_coords[0]
                        x2, y2 = line_coords[1]

                        # 查找最近点
                        nearest_points = self.find_nearest_points(
                            x1, y1, x2, y2)

                        # 显示最近点（最多显示10个）
                        for k, point in enumerate(nearest_points[:10]):
                            px, py, pz, dist = point
                            marker = self.ax.plot(
                                px,
                                py,
                                'bo',
                                markersize=8,
                                label='最近数据点'
                                if len(self.nearest_point_markers) == 0 else
                                "")[0]
                            self.nearest_point_markers.append(marker)

                            # 添加注释
                            self.ax.annotate(
                                f'({px:.1f}, {py:.1f}, {pz:.3f})', (px, py),
                                xytext=(10, 10),
                                textcoords='offset points',
                                bbox=dict(boxstyle='round,pad=0.3',
                                          facecolor='yellow',
                                          alpha=0.7),
                                fontsize=9)

                    # 清除断面图相关内容
                    for line in self.section_lines:
                        line.remove()
                    self.section_lines.clear()

                    self.ax_section.clear()
                    self.ax_section.set_title('断面图')
                    self.ax_section.set_xlabel('距离 (m)')
                    self.ax_section.set_ylabel('高程 (m)')

                    # 重新绘制图例
                    self.ax.legend()

                    # 刷新图形
                    self.fig.canvas.draw()

    def on_key_press(self, event):
        if event.key == 'control':
            self.l_key_pressed = True

    def on_key_release(self, event):
        if event.key == 'control':
            self.l_key_pressed = False

    def run(self):
        self.load_data()
        self.interpolate_data()
        self.setup_plot()

        # 连接事件
        cid_click = self.fig.canvas.mpl_connect('button_press_event',
                                                self.on_click)
        cid_key_press = self.fig.canvas.mpl_connect('key_press_event',
                                                    self.on_key_press)
        cid_key_release = self.fig.canvas.mpl_connect('key_release_event',
                                                      self.on_key_release)

        # 显示图表
        plt.tight_layout()
        plt.show()


# 主程序入口
if __name__ == "__main__":
    plotter = ContourPlotter()
    plotter.run()
